import os

from PIL.Image import Image

from Cover import Cover
from HousePlan import HousePlan
from base.ocr_base_tool import convert_pdf_to_jpg_and_collect_paths


class AddressMatcher:
    def __init__(self, input_path, output_dir='output'):
        self._input_path = input_path  # os.path.abspath(input_path)
        self.output_dir = self._input_path + r'/' + output_dir
        if not os.path.exists(self.output_dir):
            os.makedirs(self.output_dir)

    def match_addresss(self):
        try:
            # 1) 找到封面文件夹和图片
            print("1) 找到封面文件夹和图片 ###############################################################################")
            cover_dir = os.path.join(self._input_path, '封面')
            cover_images = [os.path.join(cover_dir, f) for f in os.listdir(cover_dir) if f.endswith(('jpg', 'png'))]

            # 2) 识别封面图片
            print("2) 识别封面图片 #####################################################################################")
            covers = []
            for image_path in cover_images:
                try:
                    cover = Cover(output_path=self._input_path)
                    cover.recognize_cover(image_path)
                    covers.append(cover)
                except Exception as e:
                    print(f"Error recognizing cover for image {image_path}: {e}")

            # 3) 拆分PDF为JPG
            print("3) 拆分PDF为JPG #####################################################################################")
            page_images_list = convert_pdf_to_jpg_and_collect_paths(self._input_path,
                                                                    self._input_path + r'\temp_house_plan')

            # 4) 识别分户图图片
            print("4) 识别分户图图片 #####################################################################################")
            house_plans = []
            for image_path in page_images_list:
                try:
                    house_plan = HousePlan()
                    house_plan.recognize_house_plan_on_fixed_location(image_path)
                    house_plans.append(house_plan)
                except Exception as e:
                    print(f"Error recognizing plot plan for image {image_path}: {e}")

            # 5) 匹配坐落并处理输出
            print("5) 匹配坐落并处理输出 #############################################################################")
            for cover in covers:
                for house_plan in house_plans:
                    try:
                        if cover.address == house_plan.address:
                            output_folder = os.path.join(self.output_dir, cover.registration_number)
                            if not os.path.exists(output_folder):
                                os.makedirs(output_folder)
                            output_image_path = os.path.join(output_folder, '附图0001.jpg')
                            # 假设house_plan包含获取分户图JPG文件路径的方法
                            with Image.open(house_plan.image_path) as img:
                                img.save(output_image_path)
                    except Exception as e:
                        print(f"Error matching address or saving image: {e}")
                        break  # 假设一个封面只匹配一个分户图

        except Exception as e:
            print(f"An unexpected error occurred: {e}")